# name : simu_surv_with_ia using coxph.fit and survdiff
# key: simu_surv_with_ia_fit
# contributor: Shuguang Sun
# --
library(survival)

pwe_power <- function(medCtrl = 12,
                      time.cut.p = 0,
                      hr.post = 0.69,
                      enrlBymonth = 1,
                      ratio = 2 / 3,
                      drop.rate = 0.5,
                      e.IA = 275,
                      e.FA = 335,
                      b.IA = 0.044,
                      b.FA = 0.0239,
                      hr.IA = 0.77,
                      hr.FA = 0.77,
                      hr.cnst = 0.845,
                      R = 1000,
                      seed = 520229120) {
  hRtCtrl <- -log(0.5) / medCtrl
  if (any(time.cut.p > 0)) {
    hRtExpr <- hRtCtrl * c(1, hr.post)
  } else {
    hRtExpr <- hRtCtrl * hr.post
  }
  time.cut <- c(0, time.cut.p, Inf)
  hRtDrop <- (-log(1 - drop.rate) / 12) # hazard rate of dropout, 5% per 12 months

  ## total sample size enrolled
  n <- sum(enrlBymonth)
  RANDmonth <- rep(c(0:(length(enrlBymonth) - 1)), enrlBymonth)

  HRobs.IA <- numeric(R)
  HRobs.FA <- numeric(R)
  logrank.IA <- numeric(R)
  logrank.FA <- numeric(R)
  HRsucc.IA <- numeric(R)
  HRsucc.FA <- numeric(R)
  IAsuccess <- numeric(R)
  FAsuccess <- numeric(R)
  success <- numeric(R)
  t.IA1 <- numeric(R)
  t.FA1 <- numeric(R)
  const.IA <- numeric(R)
  pchi2.IA <- numeric(R)

  for (r in 1:R) {
    # Step1: simulate data at each run
    temp <- (runif(n) <= ratio)
    RANDmonth.trt <- RANDmonth[temp == TRUE] # pop=1
    RANDmonth.ctr <- RANDmonth[temp == FALSE] # pop=2
    n.trt <- length(RANDmonth.trt)
    n.ctr <- length(RANDmonth.ctr)

    if (any(time.cut.p > 0)) {
      tte.trt <- RANDmonth.trt + gnr.time.Expr(n = n.trt, hRtExpr, time.cut)
    } else {
      tte.trt <- RANDmonth.trt + (-log(1 - runif(n.trt)) / hRtExpr)
    }
    ttdrop.trt <- RANDmonth.trt + (-log(1 - runif(n.trt)) / hRtDrop)
    tte.ctr <- RANDmonth.ctr + (-log(1 - runif(n.ctr)) / hRtCtrl)
    ttdrop.ctr <- RANDmonth.ctr + (-log(1 - runif(n.ctr)) / hRtDrop)

    tte <- c(tte.ctr, tte.trt)
    ttdrop <- c(ttdrop.ctr, ttdrop.trt)
    EVENT <- as.numeric(tte <= ttdrop)
    if (any(EVENT == 0)) {
      tte[EVENT == 0] <- ttdrop[EVENT == 0]
    }
    enrlt <- c(RANDmonth.ctr, RANDmonth.trt)
    Arm <- c(rep.int("CTR", n.ctr), rep.int("TRT", n.trt))

    pid <- order(tte)
    tte <- tte[pid]
    EVENT <- EVENT[pid]
    enrlt <- enrlt[pid]
    Arm <- Arm[pid]

    cumEVENT <- cumsum(EVENT)
    t.IA <- tte[max(which(cumEVENT <= e.IA))]
    AVAL.IA <- tte
    AVAL.IA[tte > t.IA] <- t.IA
    AVAL.IA <- AVAL.IA - enrlt
    EVENT.IA <- EVENT
    EVENT.IA[tte > t.IA] <- 0
    t.FA <- tte[max(which(cumEVENT <= e.FA))]
    AVAL.FA <- tte
    AVAL.FA[tte > t.FA] <- t.FA
    AVAL.FA <- AVAL.FA - enrlt
    EVENT.FA <- EVENT
    EVENT.FA[tte > t.FA] <- 0

    t.IA1[r] <- t.IA
    t.FA1[r] <- t.FA

    # step 2: estimate and test
    # HR obs:
    rx <- model.matrix(~trt.assign)[, -1, drop = FALSE]
    sstrata <- as.integer(strata(strataid))

    HRobs.IA[r] <- coxph.fit(
      x = rx, y = Surv(AVAL.IA, EVENT.IA),
      sstrata, offset, init, controls, weights,
      method = "efron", rownames = 1:N
    )$coefficients

    HRobs.FA[r] <- coxph.fit(
      x = rx, y = Surv(AVAL.FA, EVENT.FA),
      sstrata, offset, init, controls, weights,
      method = "efron", rownames = 1:N
    )$coefficients

    HRsucc.IA[r] <- exp(HRobs.IA[r]) < hr.IA
    HRsucc.FA[r] <- exp(HRobs.FA[r]) < hr.FA
    const.IA <- exp(HRobs.IA[r]) < hr.cnst

    # test the tteDate and check the power
    logrank.IA <- survdiff(Surv(AVAL.IA, EVENT.IA) ~ rx + strata(strataid))
    logrank.FA <- survdiff(Surv(AVAL.FA, EVENT.FA) ~ rx + strata(strataid))


    if (is.matrix(logrank.IA$obs)) {
      otmp <- apply(logrank.IA$obs, 1, sum)
      etmp <- apply(logrank.IA$exp, 1, sum)
    } else {
      otmp <- logrank.IA$obs
      etmp <- logrank.IA$exp
    }
    df.IA <- (sum(1 * (etmp > 0))) - 1

    pchi2.IA[r] <- pchisq(logrank.IA$chisq, df.IA, lower.tail = FALSE)

    if (is.matrix(logrank.FA$obs)) {
      otmp <- apply(logrank.FA$obs, 1, sum)
      etmp <- apply(logrank.FA$exp, 1, sum)
    } else {
      otmp <- logrank.FA$obs
      etmp <- logrank.FA$exp
    }
    df.FA <- (sum(1 * (etmp > 0))) - 1

    ## IAsuccess[r] <- (1-pchisq(logrank.IA$chisq,1)) < b.IA
    ## FAsuccess[r] <- !IAsuccess[r] & ((1-pchisq(logrank.FA$chisq,1)) < b.FA)
    IAsuccess[r] <- pchisq(logrank.IA$chisq, df.IA, lower.tail = FALSE) < b.IA * 2
    FAsuccess[r] <- !IAsuccess[r] & (pchisq(logrank.FA$chisq, df.FA, lower.tail = FALSE) < b.FA * 2)

    success[r] <- IAsuccess[r] | FAsuccess[r]
  }

  res <- list()
  res$HRobs.IA <- mean(exp(HRobs.IA))
  res$HRobs.FA <- mean(exp(HRobs.FA))
  res$HRsucc.IA <- mean(HRsucc.IA) * 100
  res$HRsucc.FA <- mean(HRsucc.FA) * 100
  res$const.IA <- mean(const.IA) * 100
  res$hybrid <- mean((!HRsucc.IA) & const.IA) * 100
  res$IAsuccess <- mean(IAsuccess) * 100
  res$FAsuccess <- mean(FAsuccess) * 100
  res$success <- mean(success) * 100

  porder.IA <- order(pchi2.IA)
  hrmdd_id <- max(which(pchi2.IA[porder.IA] <= b.IA))
  res$hrmdd <- exp(HRobs.IA[porder.IA][hrmdd_id])

  return(res)
}


pwe_power(
  medCtrl = 12,
  time.cut.p = 0,
  hr.post = 0.69,
  enrlBymonth = c(17, 41, 49, 35, 66, 51, 51, 38, 57, 67, 40),
  ratio = 2 / 3,
  drop.rate = 0.05,
  e.IA = 275,
  e.FA = 335,
  b.IA = 0.044,
  b.FA = 0.0239,
  hr.IA = 0.77,
  hr.FA = 0.77,
  hr.cnst = 0.845,
  R = 10000,
  seed = `(random 2147483647)`
)
